#include <stdlib.h>
#include <math.h>

#include "simulator.h"
#include "config.h"
#include "network_model_emesh_hop_counter.h"
#include "config.h"
#include "core.h"
#include "memory_manager_base.h"
#include "dvfs_manager.h"
#include "config.hpp"

//NetworkModelEMeshHopCounter::NetworkModelEMeshHopCounter(Network *net, EStaticNetwork net_type)
//   : NetworkModel(net, net_type)
//   , _hopLatency(NULL,0)
//   , _enabled(false)
//   , _num_packets(0)
//   , _num_bytes(0)
//   , _total_latency(SubsecondTime::Zero())
//{
//   SInt32 total_cores = Config::getSingleton()->getTotalCores();
//
//   _meshWidth = (SInt32) floor (sqrt(total_cores));
//   _meshHeight = (SInt32) ceil (1.0 * total_cores / _meshWidth);
//
//   try
//   {
//      _linkBandwidth = ComponentBandwidthPerCycle(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()), Sim()->getCfg()->getInt("network/emesh_hop_counter/link_bandwidth"));
//      _hopLatency = ComponentLatency(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()), Sim()->getCfg()->getInt("network/emesh_hop_counter/hop_latency"));
//
//   }
//   catch (...)
//   {
//      LOG_PRINT_ERROR("Could not read emesh_hop_counter paramters from the cfg file");
//   }
//
//   assert(total_cores <= _meshWidth * _meshHeight);
//   assert(total_cores > (_meshWidth - 1) * _meshHeight);
//   assert(total_cores > _meshWidth * (_meshHeight - 1));
//}

// zhangyan modify
NetworkModelEMeshHopCounter::NetworkModelEMeshHopCounter(Network *net, EStaticNetwork net_type)
        : NetworkModel(net, net_type)
        , _hopLatency(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()), 0)   // ✅ 传入 DVFS 参数
        , _intraChipletLatency(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()), 0)   // ✅ 片内
        , _interChipletLatency(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()), 0)   // ✅ 片间
        , _ioDieLatency(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()), 0)   // ✅ 在这里初始化
        , _ioDieBandwidth(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()), 0) // ✅ 在这里初始化
        , _enabled(false)
        , _num_packets(0)
        , _num_bytes(0)
        , _total_latency(SubsecondTime::Zero())  // SubsecondTime 是正确的
        , _intra_chiplet_bytes(0) // 统计片内通信量
        , _inter_chiplet_bytes(0) // 统计片间通信量
{
    SInt32 total_cores = Config::getSingleton()->getTotalCores();

    _meshWidth = (SInt32) floor (sqrt(total_cores));
    _meshHeight = (SInt32) ceil (1.0 * total_cores / _meshWidth);

    try
    {
        // 片内通信
        _intraChipletBandwidth = ComponentBandwidthPerCycle(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()),
                                                            Sim()->getCfg()->getInt("network/emesh_hop_counter/link_bandwidth"));
        _intraChipletLatency = ComponentLatency(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()),
                                                Sim()->getCfg()->getInt("network/emesh_hop_counter/hop_latency"));

        // 片间通信
        _interChipletBandwidth = ComponentBandwidthPerCycle(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()),
                                                            Sim()->getCfg()->getInt("network/inter_chiplet/link_bandwidth"));
        _interChipletLatency = ComponentLatency(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()),
                                                Sim()->getCfg()->getInt("network/inter_chiplet/hop_latency"));

        _useIoDie = Sim()->getCfg()->hasKey("network/io_die");  // 检查是否启用了 I/O Die

        if (_useIoDie)
        {
            // 在构造函数中读取 Sniper 配置：
            _ioDieLatency = ComponentLatency(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()),
                                             Sim()->getCfg()->getInt("network/io_die/latency"));

            _ioDieBandwidth = ComponentBandwidthPerCycle(Sim()->getDvfsManager()->getCoreDomain(net->getCore()->getId()),
                                                         Sim()->getCfg()->getInt("network/io_die/bandwidth"));
        }


        // 读取芯粒结构
        _totalCores = Sim()->getCfg()->getInt("network/chiplet/total_cores");
        _coresPerChiplet = Sim()->getCfg()->getInt("network/chiplet/cores_per_chiplet");
        _chipletWidth = Sim()->getCfg()->getInt("network/chiplet/chiplet_width");
        _chipletHeight = Sim()->getCfg()->getInt("network/chiplet/chiplet_height");
        _numChipletsX = Sim()->getCfg()->getInt("network/chiplet/num_chiplets_x");
        _numChipletsY = Sim()->getCfg()->getInt("network/chiplet/num_chiplets_y");

        String topology_str = Sim()->getCfg()->getString("network/chiplet/inter_chiplet_topology");
        std::string topology = std::string(topology_str.c_str());

        if (topology == "mesh")
            _interChipletTopology = MESH;
        else if (topology == "bus")
            _interChipletTopology = BUS;
        else if (topology == "hybrid")
            _interChipletTopology = HYBRID;
        else
            LOG_PRINT_ERROR("Invalid inter_chiplet_topology: %s", topology.c_str());
    }
    catch (...)
    {
        LOG_PRINT_ERROR("Could not read chiplet parameters from the cfg file");
    }


    assert(total_cores <= _meshWidth * _meshHeight);
    assert(total_cores > (_meshWidth - 1) * _meshHeight);
    assert(total_cores > _meshWidth * (_meshHeight - 1));

    // 通信量统计
    core_id_t core_id = getNetwork()->getCore()->getId();
    std::string netName = std::string("network.") + EStaticNetworkStrings[net_type];
    registerStatsMetric(netName.c_str(), core_id, "intra_chiplet_bytes", &_intra_chiplet_bytes);
    registerStatsMetric(netName.c_str(), core_id, "inter_chiplet_bytes", &_inter_chiplet_bytes);
}

NetworkModelEMeshHopCounter::~NetworkModelEMeshHopCounter()
{
}

// zhangyan modify
SInt32 NetworkModelEMeshHopCounter::computeChipletID(core_id_t core)
{
    SInt32 x, y;
    computePosition(core, x, y);

    SInt32 chiplet_x = x / _chipletWidth;
    SInt32 chiplet_y = y / _chipletHeight;

    return chiplet_x + chiplet_y * _numChipletsX;
}

void NetworkModelEMeshHopCounter::computePosition(core_id_t core,
                                             SInt32 &x, SInt32 &y)
{
   x = core % _meshWidth;
   y = core / _meshWidth;
}

SInt32 NetworkModelEMeshHopCounter::computeDistance(SInt32 x1, SInt32 y1, SInt32 x2, SInt32 y2)
{
   return abs(x1 - x2) + abs(y1 - y2);
}

//void NetworkModelEMeshHopCounter::routePacket(const NetPacket &pkt,
//                                         std::vector<Hop> &nextHops)
//{
//   UInt32 pkt_length = getNetwork()->getModeledLength(pkt);
//
//   SubsecondTime serialization_latency = computeSerializationLatency(pkt_length);
//
//   SInt32 sx, sy, dx, dy;
//
//   computePosition(pkt.sender, sx, sy);
//
//   if (pkt.receiver == NetPacket::BROADCAST)
//   {
//      UInt32 total_cores = Config::getSingleton()->getTotalCores();
//
//      SubsecondTime curr_time = pkt.time;
//      // There's no broadcast tree here, but I guess that won't be a
//      // bottleneck at all since there's no contention
//      for (SInt32 i = 0; i < (SInt32) total_cores; i++)
//      {
//         computePosition(i, dx, dy);
//
//         SubsecondTime latency = computeDistance(sx, sy, dx, dy) * _hopLatency.getLatency();
//         if (i != pkt.sender)
//            latency += serialization_latency;
//
//         Hop h;
//         h.final_dest = i;
//         h.next_dest = i;
//         h.time = curr_time + latency;
//
//         // Update curr_time
//         curr_time += serialization_latency;
//
//         nextHops.push_back(h);
//      }
//   }
//   else
//   {
//      computePosition(pkt.receiver, dx, dy);
//
//      SubsecondTime latency = computeDistance(sx, sy, dx, dy) * _hopLatency.getLatency();
//      if (pkt.receiver != pkt.sender)
//         latency += serialization_latency;
//
//      Hop h;
//      h.final_dest = pkt.receiver;
//      h.next_dest = pkt.receiver;
//      h.time = pkt.time + latency;
//
//      nextHops.push_back(h);
//   }
//}

//zhangyan modify
void NetworkModelEMeshHopCounter::routePacket(const NetPacket &pkt, std::vector<Hop> &nextHops)
{
    UInt32 pkt_length = getNetwork()->getModeledLength(pkt);
    SubsecondTime serialization_latency;

    SInt32 sx, sy, dx, dy;
    computePosition(pkt.sender, sx, sy);
    computePosition(pkt.receiver, dx, dy);

    SInt32 src_chiplet = computeChipletID(pkt.sender);
    SInt32 dst_chiplet = computeChipletID(pkt.receiver);

    // 统计片内片间通信量
    if (src_chiplet == dst_chiplet)
        _intra_chiplet_bytes += pkt_length;
    else
        _inter_chiplet_bytes += pkt_length;

    SubsecondTime latency;

    if (src_chiplet == dst_chiplet)
    {
        // 片内通信
        latency = SubsecondTime(computeDistance(sx, sy, dx, dy) * _intraChipletLatency.getLatency());
        serialization_latency = SubsecondTime(_intraChipletBandwidth.getRoundedLatency(pkt_length));
    }
    else
    {
        if (_useIoDie)
        {
            // 片间通信强制经过 I/O Die
            latency = _ioDieLatency.getLatency() * 2;  // 进出 I/O Die 各一次
            serialization_latency = _ioDieBandwidth.getRoundedLatency(pkt_length);
        } else
        {
            // 片间通信
            if (_interChipletTopology == MESH)
            {
                SInt32 hop_count = computeDistance(src_chiplet % _numChipletsX, src_chiplet / _numChipletsX,
                                                   dst_chiplet % _numChipletsX, dst_chiplet / _numChipletsX);
                latency = hop_count * _interChipletLatency.getLatency();
            }
            else if (_interChipletTopology == BUS)
            {
                latency = _interChipletLatency.getLatency();
            }
            else if (_interChipletTopology == HYBRID)
            {
                latency = (_chipletWidth + _chipletHeight) * _interChipletLatency.getLatency() / 2;
            }
            serialization_latency = _interChipletBandwidth.getRoundedLatency(pkt_length);
        }
    }

    if (pkt.receiver != pkt.sender)
        latency += serialization_latency;

    Hop h;
    h.final_dest = pkt.receiver;
    h.next_dest = pkt.receiver;
    h.time = pkt.time + latency;

    nextHops.push_back(h);
}

void
NetworkModelEMeshHopCounter::processReceivedPacket(NetPacket &pkt)
{
   ScopedLock sl(_lock);

   UInt32 pkt_length = getNetwork()->getModeledLength(pkt);

   core_id_t requester = INVALID_CORE_ID;

   if (pkt.type == SHARED_MEM_1)
      requester = getNetwork()->getCore()->getMemoryManager()->getShmemRequester(pkt.data);
   else // Other Packet types
      requester = pkt.sender;

   LOG_ASSERT_ERROR((requester >= 0) && (requester < (core_id_t) Config::getSingleton()->getTotalCores()),
         "requester(%i)", requester);

   if ( (!_enabled) || (requester >= (core_id_t) Config::getSingleton()->getApplicationCores()) )
      return;

   // LOG_ASSERT_ERROR(pkt.start_time > 0, "start_time(%llu)", pkt.start_time);

   SubsecondTime latency = pkt.time - pkt.start_time;

   _num_packets ++;
   _num_bytes += pkt_length;
   _total_latency += latency;
}

SubsecondTime
NetworkModelEMeshHopCounter::computeSerializationLatency(UInt32 pkt_length)
{
   // Send: (pkt_length * 8) bits
   // Bandwidth: (m_link_bandwidth) bits/cycle
   UInt32 num_bits = pkt_length * 8;
   return _linkBandwidth.getRoundedLatency(num_bits);
}
